import os
import cupy as cp
import numpy as np
import matplotlib.pyplot as plt
import cv2
import concurrent.futures
import time


python_file = os.path.dirname(__file__)
dataset_path = os.path.join(python_file, 'detect')
masks_path = os.path.join(dataset_path, 'mask')
pics_path = os.path.join(dataset_path, 'original')
results_path = os.path.join(dataset_path, 'res')

def read_mask(path):
    mask = cv2.imread(path)
    mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
    ret, mask = cv2.threshold(mask, 0, 1, cv2.THRESH_BINARY)
    contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    area = cp.array([cv2.contourArea(_) for _ in contours])
    max_area_id = cp.argmax(area)
    for _ in range(len(contours)):
        if _ != max_area_id:
            cv2.fillPoly(mask, [contours[_]], 0)
    return cp.array(mask, dtype=float)

def get_var(mask):
    sumx = cp.sum(mask, axis=0)
    indexes = cp.argwhere(sumx != 0)
    xr = cp.min(indexes)
    xl = cp.max(indexes)
    a = int((xr - xl) * 0.5)
    dx = int((xr + xl) * 0.5)
    axisy = mask[:, dx]
    indexes = cp.argwhere(axisy != 0)
    yu = cp.min(indexes)
    yd = cp.max(indexes)
    b = int((yd - yu) * 0.5)
    dy = int((yd + yu) * 0.5)
    return [a, b, dx, dy]


def iou(mask, variables):
    union = 0.
    inter = 0.
    [a, b, dx, dy, k] = variables
    [ymax, xmax] = cp.shape(mask)
    
    x = cp.arange(xmax)
    y = cp.arange(ymax)
    x, y =cp.meshgrid(x, y)

    egg = cp.square((x - dx)/a) + cp.square((y - dy)/b) * cp.exp(k * (x - dx))
    egg = cp.where(egg <= 1, 1, 0)
    inter = egg * mask
    inter_ = cp.sum(inter)
    union = egg + mask
    union = cp.where(union != 0, 1, 0)
    union_ = cp.sum(union)

    '''
    cv2.imshow('result', cp.hstack((cp.vstack((egg, mask)), cp.vstack((inter, union)))))
    cv2.waitKey()
    cv2.destroyAllWindows()
    '''
    return inter_ / union_

def plotegg(variables):
    [a, b, dx, dy, k] = variables
    X = np.linspace(-a, a, 1000)
    Y = np.sqrt((1 - np.square(X/a)) / np.exp(k * X)) * b
    plt.plot(X+dx,Y+dy,color = 'yellow')
    plt.plot(X+dx,-Y+dy,color = 'yellow')
    plt.axis('off')

def get_egg(mask_path, pic_path, save_path):
    mask = read_mask(mask_path)
    pic =  np.array(cv2.cvtColor(cv2.imread(pic_path), cv2.COLOR_BGR2RGB))

    variables = get_var(mask)
    variables.append(0.2/30)
    variables = cp.array(variables)
    learn_rate = 0.3
    grad = 0.1/600
    iou_0 = iou(mask, variables)
    iou_p = iou_0
    iou_his = [0]
    k_his = [0]

    while True:
        
        variables[4] += learn_rate*grad
        iou_p = iou(mask, variables)
        
        if iou_p > iou_0:
            grad = iou_p - iou_0
        else:
            variables[4] -= 1.25*learn_rate*grad/iou_0
            iou_p = iou(mask, variables)
            grad = iou_p - iou_0
        iou_0 = iou_p
        iou_his.append(iou_0)
        k_his.append(variables[4])
        if iou_his[-1] <= iou_his[-2]:
            variables[4] = k_his[-2]
            break
        if iou_0 > 0.9:
            break
    if variables[4] <=0:
        num = pic_path.split("al\\")[1].split(".jpg")[0]
        print('Error k in ' + num)
        with open("error pic list.txt","a") as f:
            f.write(num+"\n") 
    plt.cla()
    plt.imshow(pic)
    plotegg(cp.asnumpy(variables))
    plt.savefig(save_path, bbox_inches='tight', pad_inches=0)

def main(pic):
    
    mask_path = os.path.join(masks_path, pic)
    pic_path = os.path.join(pics_path, pic)
    result_path = os.path.join(results_path, pic)
    
    try:
        get_egg(mask_path, pic_path, result_path)
    except Exception as e:
        print('\nERROR:%s'%pic)
        print(repr(e))
    


if __name__ == '__main__':
    with open("error pic list.txt","w") as f:
            f.write("")  
    
    start = time.time()
    pics = os.listdir(pics_path)
    with concurrent.futures.ProcessPoolExecutor() as executor:
        executor.map(main, pics)
    end = time.time()
    print('TIME: ', end - start)